import cv2
import os
import numpy as np
import os.path as osp

depth_dir = "/cv/yc/DSGN2/img_bbox3d/ww_val_bev_stereo1019/pred_depths_masks"
img_dir = "/cv/yc/DSGN2/data/ww/training/image_2"
out_dir = osp.join(osp.dirname(depth_dir), "fusion")
if not os.path.exists(out_dir):
    os.makedirs(out_dir)
files = os.listdir(depth_dir)
for f in files:
    img_file = f.replace("_pred_depth", "")
    depth_img = cv2.imread(osp.join(depth_dir, f), -1)
    depth_img_new = cv2.applyColorMap(cv2.convertScaleAbs(depth_img, alpha=3), cv2.COLORMAP_JET)
    img_path = osp.join(img_dir, img_file)
    img = cv2.imread(img_path)
    old_h, old_w = img.shape[:2]
    crop_h, crop_w = min(480, old_h), min(928, old_w)
    x1 = int((old_w - crop_w) * 0.5)
    y1 = int((old_h - crop_h) * 1)
    y2 = y1 + crop_h
    img_new = img[y1:y2, x1:x1+crop_w]
    depth_img_new = depth_img_new[:, x1:x1+crop_w]
    # img_concat = np.hstack([img_new, depth_img_new])
    # import pdb
    # pdb.set_trace()
    img_concat = cv2.addWeighted(img_new, 0.5, depth_img_new, 0.5, 0)
    out_img_path = osp.join(out_dir, f)
    cv2.imwrite(out_img_path, img_concat)